JestによるAWS Lambda関数のテスト時に環境変数の設定値を一時的に変更する

こんにちは、CX事業本部 IoT事業部の若槻です。

デプロイ済みのAWS Lambda関数に対するテスト（Jestを利用）の実施時に、Lambda関数が使用する環境変数の設定値をテスト用の値に変更したいケースがあったので、対処した内容をご紹介します。

環境構築

まずテスト対象のLambda関数を作ります。

Lambda関数のリソースをCDKで定義します。関数が使用する環境変数HOGE_ENVをここで設定しています。

import { aws_lambda_nodejs, Stack, StackProps } from 'aws-cdk-lib';
import { Construct } from 'constructs';

export class AwsCdkAppStack extends Stack {
  constructor(scope: Construct, id: string, props?: StackProps) {
    super(scope, id, props);

    new aws_lambda_nodejs.NodejsFunction(this, 'hogeFunc', {
      functionName: 'hogeFunc',
      entry: 'src/hoge.ts',
      environment: {
        HOGE_ENV: 'valueHoge',
      },
    });
  }
}

Lambda関数のハンドラーでは環境変数HOGE_ENVを取得してログ出力するようにしています。（デモ用なので最低限の実装です）

const HOGE_ENV = process.env.HOGE_ENV || '';

export const handler = async (): Promise<void> => {
  //some process using environment variables

  console.log(HOGE_ENV);
};

CDK DeployしてLambda関数を作成します。

作成されたLambda関数を実行するとvalueHogeが取得されています。

テスト

Jestのテストコードは次のようになります。

import {
  GetFunctionConfigurationCommand,
  InvokeCommand,
  LambdaClient,
  UpdateFunctionConfigurationCommand,
} from '@aws-sdk/client-lambda';

const lambdaClient = new LambdaClient({ region: 'ap-northeast-1' });

beforeAll(async () => {
  //hogeFuncのConfigurationを取得
  const hogeFuncConfiguration = await lambdaClient.send(
    new GetFunctionConfigurationCommand({
      FunctionName: 'hogeFunc',
    })
  );

  //hogeFuncの環境変数をローカルの環境変数にバックアップ
  if (
    hogeFuncConfiguration.Environment &&
    hogeFuncConfiguration.Environment.Variables
  ) {
    process.env.HOGEFUNC_ORIGINAL_VARIABLES = JSON.stringify(
      hogeFuncConfiguration.Environment.Variables
    );
  } else {
    throw new Error('Failed to retrieve hogeFunc enviromnet variables.');
  }
});

afterAll(async () => {
  //ローカルの環境変数からhogeFuncの環境変数をリストア
  const HOGEFUNC_ORIGINAL_VARIABLES = process.env.HOGEFUNC_ORIGINAL_VARIABLES;
  if (HOGEFUNC_ORIGINAL_VARIABLES) {
    await lambdaClient.send(
      new UpdateFunctionConfigurationCommand({
        FunctionName: 'hogeFunc',
        Environment: { Variables: JSON.parse(HOGEFUNC_ORIGINAL_VARIABLES) },
      })
    );
  } else {
    throw new Error('Failed to restore hogeFunc enviromnet variables.');
  }
});

test('hogeFunc test', async () => {
  //hogeFuncの環境変数HOGE_ENVにテスト用の値を設定
  const HOGEFUNC_ORIGINAL_VARIABLES =
    process.env.HOGEFUNC_ORIGINAL_VARIABLES || '';
  await lambdaClient.send(
    new UpdateFunctionConfigurationCommand({
      FunctionName: 'hogeFunc',
      Environment: {
        Variables: {
          ...JSON.parse(HOGEFUNC_ORIGINAL_VARIABLES),
          HOGE_ENV: 'valueHogeForTest',
        },
      },
    })
  );

  await new Promise((r) => setTimeout(r, 1000)); //反映完了まで待機

  //環境変数の更新が反映されていることの確認のため、hogeFuncをテスト実行
  await lambdaClient.send(
    new InvokeCommand({
      FunctionName: 'hogeFunc',
    })
  );

  //some test
});

• beforeAllによる前処理で、オリジナルの環境変数の設定値をローカルの環境変数にバックアップしています。
• afterAllによる後処理で、ローカルの環境変数の値をLambda関数にリストアしています。
• testでは、冒頭でテストに使用したい設定値に環境変数を更新しています。

テストを実行します。

$ npx jest test/hogeFunc.test.ts
 PASS  test/aws-cdk-app.test.ts
  ✓ hogeFunc test (1297 ms)

Test Suites: 1 passed, 1 total
Tests:       1 passed, 1 total
Snapshots:   0 total
Time:        4.45 s, estimated 5 s
Ran all test suites matching /test\/aws-cdk-app.test.ts/i.

Lambda関数の出力をコンソールから確認すると、テスト用の値valueHogeForTestがちゃんと設定できています！

またLambda関数の環境変数の設定状況を確認すると、元の値valueHogeにちゃんとリストアできています！

テスト内でWaitを入れない場合

テストコードではLambda関数の環境変数をテスト用の値に更新した後に1秒間のWaitを入れていますが、Waitをしない場合は更新が反映される前にLambdaが実行される可能性があります。また更新反映が完了していない状態でafterAll()による環境変数の復旧処理が行われると、次のようなResourceConflictExceptionが発生します。

$ npx jest test/hogeFunc.test.ts
 FAIL  test/aws-cdk-app.test.ts


  ● Test suite failed to run

    ResourceConflictException: The operation cannot be performed at this time. An update is in progress for resource: arn:aws:lambda:ap-northeast-1:XXXXXXXXXXXX:function:hogeFunc

      33 |   const HOGEFUNC_ORIGINAL_VARIABLES = process.env.HOGEFUNC_ORIGINAL_VARIABLES;
      34 |   if (HOGEFUNC_ORIGINAL_VARIABLES) {
    > 35 |     await lambdaClient.send(
         |     ^
      36 |       new UpdateFunctionConfigurationCommand({
      37 |         FunctionName: 'hogeFunc',
      38 |         Environment: { Variables: JSON.parse(HOGEFUNC_ORIGINAL_VARIABLES) },

      at deserializeAws_restJson1ResourceConflictExceptionResponse (node_modules/@aws-sdk/client-lambda/dist-cjs/protocols/Aws_restJson1.js:5893:23)
      at deserializeAws_restJson1UpdateFunctionConfigurationCommandError (node_modules/@aws-sdk/client-lambda/dist-cjs/protocols/Aws_restJson1.js:5351:25)
      at node_modules/@aws-sdk/middleware-serde/dist-cjs/deserializerMiddleware.js:7:24
      at node_modules/@aws-sdk/middleware-signing/dist-cjs/middleware.js:11:20
      at StandardRetryStrategy.retry (node_modules/@aws-sdk/middleware-retry/dist-cjs/StandardRetryStrategy.js:51:46)
      at node_modules/@aws-sdk/middleware-logger/dist-cjs/loggerMiddleware.js:6:22
      at test/aws-cdk-app.test.ts:35:5

Lambda関数のConfigurationを更新する際は、必要に応じてWaitを入れるようにしましょう。

おわりに

JestによるAWS Lambda関数のテスト時に環境変数の設定値を一時的に変更してみました。

テスト用の値を設定するかどうかは、環境変数から取得した値をどういう用途に使用するかにもよると思います。私のケースの場合は設定値から取得したAPIエンドポイントをLambda関数が叩いてデータを取得するということを行っており、その際に取得したデータに応じたテストケースを作りたかったため、今回のような方法を取りました。（なので実際にはダミーのデータを返すREST APIを別途作ったりしています）

以上